package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{TfIdfEncoder, TfIdfEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite

/**
  * @Author songhaicheng
  * @Date 2018/08/30
  * @Email: haicheng_song@sui.com
  */
class TfIdfEncoderSuite extends FunSuite {

  val tfidf = new TfIdfEncoder()
  val paramStr: String =
    """
      |{
      | "inputCol": "feature",
      |	"preserveCols": "label",
      | "numFeatures": 20
      |}
    """.stripMargin
  val param: TfIdfEncoderParam = tfidf.parseParam(new ToolClass().encrypt(paramStr))

  test("tf-idf parameter construction test") {
    assert(param.input_pt == null && param.output_pt == null && param.modelPath == null &&
      "feature".equals(param.inputCol) && "features".equals(param.outputCol) &&
      "label".equals(param.preserveCols) && param.numFeatures == 20 && param.minDocFreq == 0)
  }

  test("tf-idf transforming function test") {
    val rawData = TestingDataGenerator.make2ColData(
      Seq(
        ("bill1", "aaa,bbb,ccc,ddd"),
        ("bill1", "aaa,ddd,eee,fff,ggg,hhh"),
        ("bill2", "aaa,ddd,ggg,iii,jjj,kkk"),
        ("bill2", "bbb,fff,iii"),
        ("bill2", "bbb,ccc,ddd,ggg,kkk"),
        ("bill3", "ccc,eee,fff,ggg,kkk"),
        ("bill3", "bbb,ccc,ddd,ggg")
      ), "label", "feature"
    )

    val tfidfData = tfidf.process(param, rawData)
    assert(tfidfData.schema.fieldNames.length == 2 && tfidfData.count() == 7)
  }

}
